import torch
from torch import nn as nn
import torch_npu
import torch.nn.functional as F
import random
import bson
import os


class GEGLU(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x: torch.Tensor):
        a, b = x.chunk(2, dim=-1)
        return a * F.gelu(b)


if __name__ == "__main__":
    op = GEGLU()
    shape_lists = [(1280,), (8, 256, 8, 128),...]
    range_lists =  [[0, 1], ...]
    dtype_lists = [torch.float16, torch.float32,...]
    
    num_cases = 50

    for i in range(num_cases):
        shape = random.choice(shape_lists)
        range_ = random.choice(range_lists)
        dtype = random.choice(dtype_lists)
        x = torch.randn(size=shape, dtype=dtype)
        x.uniform_(*range_)
        y = op(x)
        id_ = str(bson.ObjectId())
        inputs = {"op": "geglu_op", "parameters":[x], "id": id_}
        torch.save(inputs, f"/cache/data/geglu_{i}.pt")
        torch.save(y.cpu(), f"/cache/gt/geglu_{i}.pt")
